import threading
from Crypto.Cipher import AES, DES
from binascii import b2a_hex, a2b_hex
from gmssl import sm4
import base64

SECRET_KEY = b"E7lME9qbwfc="
KEY_MAGIC = b"njzf_dsas_240612"
SM4_BASE64_KEY = "Tmp6ZjE5ODQhKCokISEhIQ=="


def synchronized(func):
    func.__lock__ = threading.Lock()

    def synced_func(*args, **kwargs):
        with func.__lock__:
            return func(*args, **kwargs)

    return synced_func


class UCrypto:
    """
    加密类
    """
    __aes_mode = AES.MODE_CBC
    __des_mode = DES.MODE_CBC

    @staticmethod
    def text_encoding(origin_data) -> bytes:
        """
        数据填充，16的整数倍
        :param origin_data: 待填充的字节数据对象
        """
        if isinstance(origin_data, str):
            data = origin_data.encode()
        elif isinstance(origin_data, bytes):
            data = origin_data
        else:
            return b""

        pad_len = 16 - len(data) % 16
        return data + bytes([pad_len] * pad_len)

    @staticmethod
    def bytes_decoding(data: bytes) -> str:
        """
        数据去除填充，与pad相反
        :param data: 待去除填充的字节数据对象
        """
        pad_len = data[-1]
        return data[:-pad_len].decode()

    @classmethod
    def encrypt(cls, plain_text, cipher_encoding: str = "", encrypt_mode: str = "AES") -> [bytes, str]:
        """
        加密函数，如果text不足16位就用空格补足为16位，
        如果大于16, 不是16的倍数，那就补足为16的倍数。
        :param plain_text: 待加密文本, str or bytes类型
        :param encrypt_mode: 加密模式，支持AES算法加密和DES算法加密
        :param cipher_encoding: 密文编码格式，默认纯字节流，支持base64格式，ascii格式
        :return:
        """
        if not encrypt_mode:
            return plain_text

        if encrypt_mode == "DES":
            return cls.des_encrypt(plain_text, cipher_encoding=cipher_encoding)
        elif encrypt_mode == "SM4":
            return cls.sm4_encrypt(plain_text, cipher_encoding=cipher_encoding)
        else:
            return cls.aes_encrypt(plain_text, cipher_encoding=cipher_encoding)

    @classmethod
    def decrypt(cls, cipher_text, cipher_encoding: str = "", encrypt_mode: str = "AES") -> [bytes, str]:
        """
        加密函数，如果text不足16位就用空格补足为16位，
        如果大于16, 不是16的倍数，那就补足为16的倍数。
        :param cipher_text: 密文
        :param encrypt_mode: 加密模式，支持AES算法加密和DES算法加密
        :param cipher_encoding: 密文编码格式，默认纯字节流，支持base64格式，ascii格式
        :return:
        """
        if not encrypt_mode:
            return cipher_text

        if encrypt_mode == "DES":
            return cls.des_decrypt(cipher_text, cipher_encoding=cipher_encoding)
        elif encrypt_mode == "SM4":
            return cls.sm4_decrypt(cipher_text, cipher_encoding=cipher_encoding)
        else:
            return cls.aes_decrypt(cipher_text, cipher_encoding=cipher_encoding)

    @classmethod
    def sm4_encrypt(cls, plain_text, secret_key: str = SM4_BASE64_KEY, cipher_encoding: str = "") -> bytes:
        """
        国密4加密函数
        :param plain_text: 待加密文本
        :param secret_key: 密钥
        :param cipher_encoding: 密文编码格式，默认纯字节流，支持base64格式
        :return:
        """
        sm4_encrypt = sm4.CryptSM4()
        sm4_encrypt.set_key(base64.b64decode(secret_key), sm4.SM4_ENCRYPT)
        cipher_text = sm4_encrypt.crypt_ecb(plain_text.encode())

        if cipher_encoding == "base64":
            return base64.b64encode(cipher_text)
        return cipher_text

    @classmethod
    def sm4_decrypt(cls, cipher_text, secret_key: str = SM4_BASE64_KEY, cipher_encoding: str = "") -> str:
        """
        国密4解密函数
        :param cipher_text: 加密后的base64编码数据
        :param secret_key: 密钥
        :param cipher_encoding: 密文编码格式，默认纯字节流，支持base64格式
        :return:
        """
        sm4_decrypt = sm4.CryptSM4()
        sm4_decrypt.set_key(base64.b64decode(secret_key), sm4.SM4_DECRYPT)

        if cipher_encoding == "base64":
            plain_text = sm4_decrypt.crypt_ecb(base64.b64decode(cipher_text))
        else:
            plain_text = sm4_decrypt.crypt_ecb(cipher_text)
        return plain_text.decode()

    @classmethod
    def aes_encrypt(cls, plain_text, secret_key: bytes = KEY_MAGIC, crypto_iv: bytes = b"0000000000000000",
                    cipher_encoding: str = "") -> bytes:
        """
        加密函数，如果text不足16位就用空格补足为16位，
        如果大于16, 不是16的倍数，那就补足为16的倍数。
        :param plain_text: 待加密文本, str or bytes类型
        :param secret_key: 密钥
        :param crypto_iv: 初始化向量
        :param cipher_encoding: 密文编码格式，默认纯字节流，支持base64格式，ascii格式
        :return:
        """
        crypto_inst = AES.new(secret_key, cls.__aes_mode, crypto_iv)
        text_bytes = cls.text_encoding(plain_text)

        cipher_text = crypto_inst.encrypt(text_bytes)
        if cipher_encoding == "ascii":
            return b2a_hex(cipher_text)
        elif cipher_encoding == "base64":
            return base64.b64encode(cipher_text)
        return cipher_text

    @classmethod
    def aes_decrypt(cls, cipher_text, secret_key: bytes = KEY_MAGIC, crypto_iv: bytes = b"0000000000000000",
                    cipher_encoding: str = "") -> str:
        """
        解密函数
        :param cipher_text: 加密后的base64编码数据
        :param secret_key: 密钥
        :param crypto_iv: 初始化向量
        :param cipher_encoding: 密文编码格式，默认纯字节流，支持base64格式，ascii格式
        :return:
        """
        own_crypto = AES.new(secret_key, cls.__aes_mode, crypto_iv)
        if cipher_encoding == "base64":
            plain_text = own_crypto.decrypt(base64.b64decode(cipher_text))
        elif cipher_encoding == "ascii":
            plain_text = own_crypto.decrypt(a2b_hex(cipher_text))
        else:
            plain_text = own_crypto.decrypt(cipher_text)
        return cls.bytes_decoding(plain_text)

    @classmethod
    def des_encrypt(cls, plain_text, secret_key: bytes = SECRET_KEY, crypto_iv: bytes = b"12345678",
                    cipher_encoding: str = "") -> bytes:
        """
        DES加密，返回加密后的base64编码数据
        :param plain_text: 待加密文本, str or bytes类型
        :param secret_key: 密钥
        :param crypto_iv: 初始化向量
        :param cipher_encoding: 密文编码格式，默认纯字节流，支持base64格式，ascii格式
        """
        secret_key = base64.b64decode(secret_key)
        plain_text = cls.text_encoding(plain_text)
        cipher = DES.new(secret_key, cls.__des_mode, crypto_iv)
        cipher_text = cipher.encrypt(plain_text)
        if cipher_encoding == "ascii":
            return b2a_hex(cipher_text)
        elif cipher_encoding == "base64":
            return base64.b64encode(cipher_text)
        return cipher_text

    @classmethod
    def des_decrypt(cls, cipher_text: bytes, secret_key: bytes = SECRET_KEY, crypto_iv: bytes = b"12345678",
                    cipher_encoding: str = "") -> str:
        """
        DES解密
        :param cipher_text: 加密后的base64编码数据
        :param secret_key: 密钥
        :param crypto_iv: 初始化向量
        :param cipher_encoding: 密文编码格式，默认纯字节流，支持base64格式，ascii格式
        """
        try:
            secret_key = base64.b64decode(secret_key)
            cipher = DES.new(secret_key, cls.__des_mode, crypto_iv)
            if cipher_encoding == "base64":
                plain_text = cipher.decrypt(base64.b64decode(cipher_text))
            elif cipher_encoding == "ascii":
                plain_text = cipher.decrypt(a2b_hex(cipher_text))
            else:
                plain_text = cipher.decrypt(cipher_text)
            return cls.bytes_decoding(plain_text)
        except Exception as e:
            print(e)


if __name__ == "__main__":
    crypto_object1 = "TRS"  # 子系统以名字为加密对象，大写，编辑器为EDITOR
    crypto_object2 = "url_violation_risk.py"  # celery worker以ip为加密对象
    result = UCrypto.encrypt(crypto_object2, cipher_encoding="ascii")
    pass_result = result.decode("utf-8")
    print(pass_result)
    print(type(pass_result))
    print(UCrypto.decrypt(pass_result.encode("utf-8"), cipher_encoding="ascii"))
    plaintext = "*_*Jnzf2022_*_"
    encrypted_text = UCrypto.encrypt(plaintext, cipher_encoding="base64", encrypt_mode="DES")
    print("Encrypted Text:", encrypted_text.decode())
    print("decrypt res: ", UCrypto.decrypt(encrypted_text, cipher_encoding="base64", encrypt_mode="DES"))
    sm4_plaintext = "API审计检测预警系统"
    encrypted_text = UCrypto.encrypt(sm4_plaintext, cipher_encoding="base64", encrypt_mode="SM4")
    print("SM4 Encrypted Text:", encrypted_text.decode())
    print("SM4 decrypt res: ", UCrypto.decrypt(encrypted_text, cipher_encoding="base64", encrypt_mode="SM4"))
